import os
import re
import matplotlib.pyplot as plt
from matplotlib.pyplot import MultipleLocator
def fill_result(result_file, result):
    epoch_list = []
    loss_list = []
    for line in result_file:
        if 'mask loss' in line:
            epoch, loss = re.findall(r'\d+\.*\d*', line)
            # print(epoch, loss)
            epoch_list.append(int(epoch))
            loss_list.append(float(loss))
    result.append(epoch_list)
    result.append(loss_list)


def loss_func_graph(train_result_path, val_result_path):
    train_result_file = open(train_result_path, 'r')
    val_result_file = open(val_result_path, 'r')
    train_result = []
    val_result = []
    fill_result(train_result_file, train_result)
    fill_result(val_result_file, val_result)
    x_major_locator = MultipleLocator(1)
    plt.plot(train_result[0], train_result[1], marker='*', label='train')
    plt.plot(val_result[0], val_result[1], marker='x', label='test')
    ax = plt.gca()
    ax.xaxis.set_major_locator(x_major_locator)
    plt.xlim(-0.5, 20.5)
    plt.legend()
    plt.savefig('/home/aistudio/logs/loss.jpg')
    train_result_file.close()
    val_result_file.close()


if __name__ == '__main__':
    loss_func_graph(os.path.join('/home/aistudio/logs', 'train.txt'),
                    os.path.join('/home/aistudio/logs', 'test.txt'))
